Nvidia Jetson NX 系列 03.搭建 pytorch 環境

Posted on Edited on In

因為需要在 nx 板子上測試不同 framework 的環境，因此要新增加裝 pytorch 的環境。原本的環境需求是這樣：

1
2
3
4
# install requirements
$ pip install pycocotools numpy opencv-python tqdm tensorboard tensorboardX pyyaml webcolors
$ pip install torch==1.4.0
$ pip install torchvision==0.5.0

但這邊有個血淚史，如果是要直接使用別人 github 上的 code，一如往常 CUDA 相關的環境，尤其和深度學習相關的一定要非常注意他的套件包的版本號。對失敗總結不感興趣的可以直接跳到成功安裝的那塊閱讀。

失敗情形總結

正常的 linux 環境下，可以直接看 pytorch 的官網[連結]給出的版本號配對，天知道我在這裡到底碰壁了幾次…簡單總結一下，在 Nvidia Xavier NX 的板子環境裡，沒辦法直接使用以下狀況：

方法一、pip install torch===x.x.x
1
2
3
$ pip install torch===1.5.0
ERROR: Could not find a version that satisfies the requirement torch===1.5.0 (from versions: 0.1.2, 0.1.2.post1, 0.1.2.post2)
ERROR: No matching distribution found for torch===1.5.0

會直接找不到套件包，無法使用。

方法二、pypi 網站[連結]或 pytorch.org 提供的 wheel 檔案

如果是一般的 linux 環境，可能是可行的吧？我只在終端 NX 上直接嘗試，不確定是否會遇到什麼其他意外。

這邊遇到一個錯誤，當時我以為只要找 wheel 版本檔名內容後面接 linux 就可以了，因此看到 pypi 提供的 wheel 就開開心心照著 pytorch 的官網連結內容告知的版本，CUDA 10.2 搭配 pytorch 1.5.1 + torchvision 0.6.1，殊不知我太傻太天真。

wheel 檔案安裝失敗。無論是直接將 wheel 檔案先下載好，還是利用 -f 導向網址 pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html 都不行喔！

補充：wheel 檔名會表示要配合的 cuda 版本、torchvision 版本、python 版本 以及 環境是 linux arm64 windows。 例如：cu102/torchvision-0.6.0-cp35-cp35m-linux_x86_64.whl

首先 Xavier NX 上，它需要的 wheel 檔格式一定要是 aarch 結尾，可以用指令 pip debug --verbose 查看該系統相容的 pip wheel 檔案名稱。

1
2
3
4
5
6
$ pip debug --verbose
  cp27-cp27mu-manylinux2014_aarch64
  cp27-cp27mu-linux_aarch64
  cp27-none-manylinux2014_aarch64
  cp27-none-linux_aarch64
  ...

再來是，Xavier 系列相容的 torch 版本，一定要記得看 Nvidia 的官網上指定的啊！NX CUDA 10.2 不能搭配上方說的 torch 1.5.1 + torchvision 0.6.1。

官網提供的安裝辦法

請往這裡走

貼文串的標題雖然是 Jetson Nano，但內文也說可以符合 NX 和 AGX。另外常常遇到問題後，NANO 的回答都還算是通用解，畢竟都是 Xavier 系列，而且 NANO 出來的時間比較久，也比較多討論串可以爬文。

簡言之，NX 上面的 CUDA 版本和 JetPack 版本有關係，像這塊 NX 搭配 JetPack 4.4 版本，也就是根據這篇文章我要安裝 PyTorch v1.5.0 以上，然後這邊我又再次踩雷了… 要自己注意是 JetPack 4.4 Developer 還是 Production，分別對應 1.5 和 1.6。

確定好版本後，點開 Installation，按照指令安裝，像我要搭配的是 torch 1.6.0 + torchvision 0.7.0。

1
2
3
4
$ wget https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
$ sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
$ pip3 install Cython
$ pip3 install numpy torch-1.6.0-cp36-cp36m-linux_aarch64.whl

這邊有神奇的一點是，不知道為何 numpy 一定要擺在 torch-1.6.0 前面，我之前一直沒裝好是因為我想說已經裝過 numpy 了，就把它拔掉了。

1
2
3
4
$ sudo apt-get install libjpeg-dev zlib1g-dev
$ git clone --branch release/0.7 https://github.com/pytorch/vision torchvision
$ cd torchvision
$ sudo python3 setup.py install

另外因為我搭配了 virtualenv，中間遇到一個問題。

1
2
3
4
5
$ sudo python3 setup.py install
Traceback (most recent call last):
  File "setup.py", line 13, in <module>
    import torch
ModuleNotFoundError: No module named 'torch'

但在虛擬環境下開啟 python import torch 是有東西的。因此我在 setup.py 裡面加上幾行印出系統執行檔位置，發現安裝時沒有正確的指向虛擬環境的執行檔。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
(det) $ python
Python 3.6.9 (default, Jul 17 2020, 12:50:27)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torchvision
>>> import sys
>>> print (sys.executable)
/home/user/.virtualenvs/det/bin/python # 虛擬環境指向的 python 執行檔

(det) $ sudo python3 setup.py install
/usr/bin/python3 # 在安裝 setup.py 時指向的執行檔位置
Traceback (most recent call last):
  File "setup.py", line 13, in <module>
    import torch
ModuleNotFoundError: No module named 'torch'

所以最後，其實強制導向到正確的位置進行執行即可。

1
(det) $ sudo /home/user/.virtualenvs/det/bin/python3 setup.py install

完成安裝後按照官網的測試指令，就可以確認是否安裝完成了歐！

Reference

JetPack SDK

Linux 編譯 shared library